import { app } from "/scripts/app.js";
import { api } from "/scripts/api.js";

console.log(` 添加训练组将`);

// 消息处理函数
function handleLoraMessage(event) {
    console.log("收到LoRA相关消息:", event);
    const data = event.detail;
    
    // 根据消息类型处理不同的逻辑
    switch(data.type) {
        case "preset_changed":
            console.log("预设改变:", data.preset);
            // 处理预设改变的逻辑
            break;
        case "vram_changed":
            console.log("显存设置改变:", data.vram);
            // 处理显存改变的逻辑
            break;
        case "training_progress":
            console.log("训练进度:", data.progress);
            // 处理训练进度更新
            break;
        default:
            console.log("未知消息类型:", data);
    }
}

// 基础Widget类
class BaseWidget {
    constructor() {
        this.value = null;
        this.element = null;
        this.node = null;
        this.name = null;
        this.type = "WIDGET";
        console.log(`[BaseWidget][constructor] 创建基础小部件`);
    }

    init(node, name, inputData) {
        console.log(`[BaseWidget][init] 开始初始化:`, {
            node: node?.type,
            name,
            inputData
        });
        
        this.node = node;
        this.name = name;
        this.value = inputData?.default || inputData?.options?.[0] || null;
        this.element = this.createElement();
        
        console.log(`[BaseWidget][init] 初始化完成:`, {
            name: this.name,
            value: this.value,
            element: this.element?.tagName
        });
        
        return this;
    }

    createElement() {
        throw new Error("Must be implemented by subclass");
    }

    getValue() {
        return this.value;
    }

    setValue(value, triggerChange = true) {
        if (this.value !== value) {
            this.value = value;
            if (triggerChange && this.node) {
                this.node.onWidgetChanged?.(this.name, this.value);
            }
        }
    }

    serialize() {
        return this.value;
    }
}

// 预设类型选择器
class PresetTypeWidget extends BaseWidget {
    constructor() {
        super();
        this.type = "combo";  // 使用 ComfyUI 的标准下拉框类型
    }

    createElement() {
        const div = document.createElement("div");
        div.style.width = "100%";
        
        const select = document.createElement("select");
        select.style.width = "100%";
        select.style.color = "var(--input-text)";
        select.style.backgroundColor = "var(--comfy-input-bg)";
        select.style.borderRadius = "8px";
        select.style.border = "2px solid var(--border-color)";
        select.style.padding = "4px 8px";
        
        const options = this.node.inputs?.[this.name]?.[1]?.options || [];
        options.forEach(option => {
            const opt = document.createElement("option");
            opt.value = option;
            opt.textContent = option;
            if (option === this.value) {
                opt.selected = true;
            }
            select.appendChild(opt);
        });
        
        select.addEventListener("change", (e) => {
            this.setValue(e.target.value);
            // 发送预设改变消息到服务器
            api.emit("lora_preset_changed", {
                node_id: this.node.id,
                preset: e.target.value
            });
        });
        
        div.appendChild(select);
        return div;
    }
}

// 显存大小选择器
class VRAMSizeWidget extends BaseWidget {
    constructor() {
        super();
        this.type = "combo";  // 使用 ComfyUI 的标准下拉框类型
    }

    createElement() {
        const div = document.createElement("div");
        div.style.width = "100%";
        
        const select = document.createElement("select");
        select.style.width = "100%";
        select.style.color = "var(--input-text)";
        select.style.backgroundColor = "var(--comfy-input-bg)";
        select.style.borderRadius = "8px";
        select.style.border = "2px solid var(--border-color)";
        select.style.padding = "4px 8px";
        
        const options = this.node.inputs?.[this.name]?.[1]?.options || ["4GB", "6GB", "8GB", "12GB", "16GB+"];
        options.forEach(option => {
            const opt = document.createElement("option");
            opt.value = option;
            opt.textContent = option;
            if (option === this.value) {
                opt.selected = true;
            }
            select.appendChild(opt);
        });
        
        select.addEventListener("change", (e) => {
            this.setValue(e.target.value);
            // 发送显存改变消息到服务器
            api.emit("lora_vram_changed", {
                node_id: this.node.id,
                vram: e.target.value
            });
        });
        
        div.appendChild(select);
        return div;
    }
}

// 注册扩展
app.registerExtension({
    name: "comfy.lora_train.widgets",
    async setup() {
        // 注册消息监听器
        api.addEventListener("lora_message", handleLoraMessage);
        console.log("LoRA训练组件已初始化，消息监听器已注册");
    },
    
    async beforeRegisterNodeDef(nodeType, nodeData, app) {
        if (nodeData.name !== "LoRAQuickPreset") {
            return;
        }
        
        console.log("[Extension] 注册节点:", {
            nodeType: nodeType?.name,
            nodeData: nodeData?.name,
            category: nodeData?.category,
            input: nodeData?.input?.required
        });

        // 注册小部件
        const widgetTypes = {
            PresetTypeWidget: PresetTypeWidget,
            VRAMSizeWidget: VRAMSizeWidget
        };

        // 添加小部件工厂函数
        nodeType.prototype.addCustomWidget = function(name, widgetType, options) {
            console.log("[Node] 添加小部件:", {
                name,
                type: widgetType,
                options
            });
            
            if (!widgetTypes[widgetType]) {
                console.error(`未知的小部件类型: ${widgetType}`);
                return null;
            }
            
            const widget = new widgetTypes[widgetType]();
            widget.init(this, name, options);
            
            if (!this.widgets) {
                this.widgets = [];
            }
            this.widgets.push(widget);
            
            return widget;
        };

        // 扩展原始的 ComfyNode
        const originalOnNodeCreated = nodeType.prototype.onNodeCreated;
        nodeType.prototype.onNodeCreated = function() {
            console.log("[Node] 节点创建:", this.type);
            
            if (originalOnNodeCreated) {
                originalOnNodeCreated.apply(this);
            }

            // 添加预设类型选择器
            this.addCustomWidget("preset_type", "PresetTypeWidget", {
                default: nodeData.input.required.preset_type[1].default,
                options: nodeData.input.required.preset_type[1].options
            });

            // 添加显存大小选择器
            this.addCustomWidget("vram_size", "VRAMSizeWidget", {
                default: nodeData.input.required.vram_size[1].default,
                options: nodeData.input.required.vram_size[1].options
            });
        };
    }
}); 